# -*- coding: utf-8 -*-
"""
Created on Mon Sep 23 12:35:59 2024
设计了拐角流相关的计算类
此模型关注局部拐角处的流动特性，属于拐角流corner flow的研究
@author: wangduo
"""
import numpy as np
import math
from geometry import Grid
from visualize import CornerFlowPlotter, SosiFlowPlotter
# embed figures in the notebook
#%matplotlib inline
#%config InlineBackend.figure_format='svg'# 让图片看起来像矢量图，更清晰

# 提取目标值直线路径上所有坐标点
def extract_line_points(X, Y, target_value, axis='x'):
    """
    Extracts the points along a specific line in the mesh grid.

    Parameters:
        X (2D array): The grid of x-coordinates.
        Y (2D array): The grid of y-coordinates.
        target_value (float): The specific值 along the axis to extract points.
        axis (str): 'x' to extract points along a line with a constant x value,
                    'y' to extract points along a line with a constant y value.

    Returns:
        x_values_on_line (2D array): The x-coordinates of the points on the line without NaN values.
        y_values_on_line (2D array): The y-coordinates of the points on the line without NaN values.
    """
    x_values_on_line, y_values_on_line = None, None

    if axis == 'x':
        # Find the index where X equals target_value (或最接近的索引)
        index = np.abs(X[0, :] - target_value).argmin()
        x_values_on_line = X[:, index].reshape(-1, 1)
        y_values_on_line = Y[:, index].reshape(-1, 1)
    elif axis == 'y':
        # Find the index where Y equals target_value (或最接近的索引)
        index = np.abs(Y[:, 0] - target_value).argmin()
        x_values_on_line = X[index, :].reshape(1, -1)
        y_values_on_line = Y[index, :].reshape(1, -1)
    else:
        raise ValueError("axis must be 'x' or 'y'")

    # Convert to 2D array and remove NaN values
    x_values_on_line = np.array(x_values_on_line)
    y_values_on_line = np.array(y_values_on_line)
    x_values_on_line = x_values_on_line[~np.isnan(x_values_on_line).any(axis=1)]
    y_values_on_line = y_values_on_line[~np.isnan(y_values_on_line).any(axis=1)]

    return x_values_on_line, y_values_on_line

# 90°拐角流类
class CornerFlow:
    """
    定义室内拐角流动（Corner Flow）的基本计算。

    该类用于计算势流的复势函数、流函数、速度分布等参数。

    这里用到的公式仅为直角拐角
    """

    def __init__(self, strength: float, corner_origin=(0.0, 0.0)):
        """
        初始化拐角流参数

        参数：
        ----------
        strength: float
            拐角流的强度，单位：m²/s。
        corner_origin: tuple (float, float)
            拐角流的原点坐标 (x, y)，默认为 (0.0, 0.0)。
        """
        self.A = strength  # 拐角流强度
        self.corner_origin = corner_origin  # 记录拐角流的原点
        self.x_cf, self.y_cf = corner_origin  # 解包坐标
        self.z_cf = complex(*corner_origin)  # 复数形式的拐角坐标

    def __repr__(self):
        """
        返回类的字符串表示，方便打印查看
        """
        return f"CornerFlow(strength={self.A}, origin={self.corner_origin})"

    def cplx_W(self, X: float, Y: float):
        """
        计算复势函数 W 及其对应的速度势 φ 和流函数 ψ

        公式：
            W(Z) = A * Z²
            φ = Re(W)   (速度势)
            ψ = Im(W)   (流函数)

        参数：
        ----------
        X: float
            x 方向坐标
        Y: float
            y 方向坐标

        返回：
        ----------
        w_ptn: complex
            复势 W(Z)
        phi: float
            速度势（φ）
        psi: float
            流函数（ψ）
        """
        Z = X + Y * 1j  # 复数表示坐标
        w_ptn = self.A * Z ** 2  # 计算复势 W
        phi = w_ptn.real  # 速度势
        psi = w_ptn.imag  # 流函数
        return w_ptn, phi, psi

    def cplx_velocity(self, X: float, Y: float):
        """
        计算速度场

        公式：
            dW/dZ = d(A*Z²)/dZ = 2 * A * Z
            u = Re(dW/dZ)  （x 方向速度分量）
            v = Im(dW/dZ)  （y 方向速度分量）

        参数：
        ----------
        X: float
            x 方向坐标
        Y: float
            y 方向坐标

        返回：
        ----------
        magV: float
            速度模（√(u² + v²)）
        u: float
            x 方向速度分量
        v: float
            y 方向速度分量
        """
        Z = X + Y * 1j  # 复数表示坐标
        dwdz = 2 * self.A * Z  # 计算 dW/dZ

        magV = abs(dwdz)  # 计算速度模
        u = dwdz.real  # x 方向速度分量
        v = -dwdz.imag  # y 方向速度分量 复速度是共轭速度，提取时需要变号

        return magV, u, v






"""____示例____"""
# Example of using the general models
if __name__ == "__main__":
    
    height = 2.5 # m
    length = 4.0 # m
    #geo_x = np.array([0.0, length])
    #geo_y = np.array([0.0, height])
    N = 50
    room_geo = Grid(height, length, N) # 网格坐标

    # CornerFlow
    U0 = 1 # m/s
    A = U0 / (2*height) # 1/s
    room_cf = CornerFlow(A, corner_origin=(2.0, -1.0))
    W_cf, phi_cf, psi_cf = room_cf.cplx_W(room_geo.X, room_geo.Y)
    V_cf, vx_cf, vy_cf= room_cf.cplx_velocity(room_geo.X, room_geo.Y)
    print("f")
    # 绘图
    #plt_cf = CornerFlowPlotter(height, length, room_cf, room_geo)
    #plt_cf.plot_psi(psi_cf, title='Stream Function')
    #plt_cf.plot_vxvy(vx_cf, vy_cf, title='velocity vector')


    '''
    # 物理参数

    # 源汇
    v_inlet = 40 # m/s
    d = 0.05 # m
    Strength_sosi = v_inlet * d
    pos_source = 4.5 * d # m 系数：4.5~6.5
    pos_sink = length - 0.8 # m
    room_sosi = RoomSourceSink(Strength_sosi, pos_source, pos_sink, height)
    W_sosi, phi_sosi, psi_sosi = room_sosi.cplx_W(room_geo.X, room_geo.Y)
    V_sosi, jd_sosi, vx_sosi, vy_sosi = room_sosi.cplx_velocity(room_geo.X, room_geo.Y)
    RR, II = room_sosi.cplx_2velovity(room_geo.X, room_geo.Y)

    #print(jd_sosi)
    p_sosi, CP_sosi,dv2 = room_sosi.pressure(v_inlet, room_geo.X, room_geo.Y)
    #print(dv2)
    #print(CP_sosi)
    #print(vy_sosi)

    # 角涡
    Gamma_vo = 0.05
    x_vortex = 0.1
    y_vortex = 0.1
    room_vo = RoomVortex(Gamma_vo, x_vortex, y_vortex, height)
    W_vo, phi_vo, psi_vo = room_vo.cplx_W(room_geo.X, room_geo.Y)
    V_vo, vx_vo, vy_vo = room_vo.cplx_velocity(room_geo.X, room_geo.Y)

    # 顶涡
    Gamma_vo2 = -2
    x_vortex2 = 3.5
    y_vortex2 = 2.0
    room_vo2 = RoomVortex(Gamma_vo, x_vortex2, y_vortex2, height)
    W_vo2, phi_vo2, psi_vo2 = room_vo2.cplx_W(room_geo.X, room_geo.Y)
    V_vo2, vx_vo2, vy_vo2 = room_vo2.cplx_velocity(room_geo.X, room_geo.Y)

    # 拐角流
    A = 2
    x_origin = 0.0
    y_origin = 0.0
    room_cf = RoomCorner(A, x_origin, y_origin)
    W_cf, phi_cf, psi_cf = room_cf.cplx_W(room_geo.X, room_geo.Y)
    V_cf, vx_cf, vy_cf= room_cf.cplx_velocity(room_geo.X, room_geo.Y)

    # 绘图
    plt_cf = CornerFlowPlot(height, length, room_cf, room_geo)
    #plt_cf.plot_psi(psi_cf, title='Stream Function')

    #plt_cf.plot_phi(phi_cf, title='Potential Function')
    plt_cf.plot_flow_net(phi_cf, psi_cf, title='Flow Net')
    plt_cf.plot_vxvy(vx_cf, vy_cf, title='Velocity Field')

    plt_cf.plot_vxvy(RR, II, title='2 order')
    plt_cf.plot_MagV(V_cf, title='Velocity Magnitude')
    #plt_cf.plot_pressure(CP_cf, title='Pressure Field')

    # 绘图（源汇SC变换模型）
    target_psi = 0.0182
    plt_sosi = RoomFlowPlot(height, length, room_sosi, room_geo)
    plt_sosi.plot_psi(psi_sosi, title='Stream Function')

    #plt_sosi.plot_phi(phi_sosi, title='Potential Function')
    #plt_sosi.plot_flow_net(phi_sosi, psi_sosi, title='Flow Net')
    #plt_sosi.plot_vxvy(vx_sosi, vy_sosi, title='Velocity Field')

    #plt_sosi.plot_vxvy(RR, II, title='2 order')
    #plt_sosi.plot_MagV(V_sosi, title='Velocity Magnitude')
    #plt_sosi.plot_pressure(CP_sosi, title='Coefficient of Pressure')
    '''
